#!/usr/bin/env python3
print("Loading modules...")
import argparse
import yaml
from DeepPurpose import DTI as models
from DeepPurpose.dataset import process_BindingDB
from DeepPurpose.utils import data_process, generate_config


def main():
    parser = argparse.ArgumentParser(
        description="Train a DTI prediction model using BindingDB data."
    )

    # Add command line arguments
    parser.add_argument(
        "--input",
        type=str,
        required=True,
        help="Path to the BindingDB DTI data file (TSV format, can be gz compressed).",
    )
    parser.add_argument(
        "--config",
        type=str,
        required=True,
        help="Path to the YAML configuration file for model training.",
    )
    parser.add_argument(
        "--output", type=str, required=True, help="Path to save the trained model."
    )
    parser.add_argument(
        "--metric",
        type=str,
        default="IC50",
        help="Metric to use for processing the data. Default is IC50.",
    )
    parser.add_argument(
        "--split_method",
        type=str,
        default="cold_protein",
        help="Method to split the data. Default is cold_protein.",
    )
    parser.add_argument(
        "--frac",
        nargs=3,
        type=float,
        default=[0.7, 0.1, 0.2],
        help="Fractions for train, validation, and test splits. Default is [0.7, 0.1, 0.2].",
    )

    args = parser.parse_args()

    db_path = args.input
    config_path = args.config
    model_path = args.output
    metric = args.metric
    split_method = args.split_method
    frac = args.frac

    # Load the YAML configuration file
    with open(config_path, "r") as f:
        yaml_config = yaml.safe_load(f)

    # Ensure drug_encoding and target_encoding are present in the config
    if "drug_encoding" not in yaml_config or "target_encoding" not in yaml_config:
        raise ValueError(
            "The configuration file must contain 'drug_encoding' and 'target_encoding' keys."
        )

    drug_encoding = yaml_config["drug_encoding"]
    target_encoding = yaml_config["target_encoding"]

    print(f"Processing BindingDB data from {db_path}...")
    X_drug, X_target, y = process_BindingDB(
        path=db_path, y=metric, binary=False, convert_to_log=True
    )

    print("Processing data...")
    train, val, test = data_process(
        X_drug,
        X_target,
        y,
        drug_encoding=drug_encoding,
        target_encoding=target_encoding,
        split_method=split_method,
        frac=frac,
    )

    yaml_config["result_folder"] = model_path

    config = generate_config(**yaml_config)
    net = models.DBTA(**config)

    print("Training the model...")
    net.train(train, val, test)

    print(f"Saving the model to {model_path}...")
    net.save_model(model_path)


if __name__ == "__main__":
    main()
